from django.core.management.base import BaseCommand
from django_celery_beat.models import PeriodicTask, CrontabSchedule
import json

class Command(BaseCommand):
    help = '设置定期数据分析任务'

    def handle(self, *args, **kwargs):
        try:
            # 创建或更新每日任务的Crontab计划
            daily_schedule, _ = CrontabSchedule.objects.get_or_create(
                hour=15,
                minute=30,
                day_of_week='*',
                day_of_month='*',
                month_of_year='*',
                timezone='Asia/Shanghai'
            )

            # 创建或更新每周任务的Crontab计划
            weekly_schedule, _ = CrontabSchedule.objects.get_or_create(
                hour=15,
                minute=30,
                day_of_week='5',  # 周五
                day_of_month='*',
                month_of_year='*',
                timezone='Asia/Shanghai'
            )

            # 创建或更新每月任务的Crontab计划（设置为每月28日，考虑到不同月份的天数）
            monthly_schedule, _ = CrontabSchedule.objects.get_or_create(
                hour=15,
                minute=30,
                day_of_week='*',
                day_of_month='28',  # 设置为每月28日
                month_of_year='*',
                timezone='Asia/Shanghai'
            )

            # 设置每日数据分析任务
            PeriodicTask.objects.update_or_create(
                name='每日数据分析',
                defaults={
                    'task': 'data_analysis.tasks.analyze_market_data',
                    'crontab': daily_schedule,
                    'args': json.dumps(['day']),
                    'enabled': True,
                }
            )

            # 设置每周数据分析任务
            PeriodicTask.objects.update_or_create(
                name='每周数据分析',
                defaults={
                    'task': 'data_analysis.tasks.analyze_market_data',
                    'crontab': weekly_schedule,
                    'args': json.dumps(['week']),
                    'enabled': True,
                }
            )

            # 设置每月数据分析任务
            PeriodicTask.objects.update_or_create(
                name='每月数据分析',
                defaults={
                    'task': 'data_analysis.tasks.analyze_market_data',
                    'crontab': monthly_schedule,
                    'args': json.dumps(['month']),
                    'enabled': True,
                }
            )

            self.stdout.write(self.style.SUCCESS('成功设置定期数据分析任务'))

        except Exception as e:
            self.stdout.write(self.style.ERROR(f'设置定期任务失败: {str(e)}'))
            raise